package com.liezh.server;

import org.apache.thrift.TException;
import org.apache.thrift.protocol.*;
import org.apache.thrift.transport.TIOStreamTransport;
import org.apache.thrift.transport.TTransport;

import java.io.BufferedInputStream;
import java.lang.reflect.Field;
import java.util.HashMap;
import java.util.Map;

/**
 * @author liezh
 * @date 2022-3-7
 */
public class THeaderServerProtocol extends TBinaryProtocol {

    private Map<String, String> HEAD_INFO;

    public THeaderServerProtocol(TTransport transport) {
        super(transport);
        HEAD_INFO = new HashMap<>();
    }

    public boolean readFieldZero() throws TException {
        TField schemeField = this.readFieldBegin();

        if (schemeField.id == 0 && schemeField.type == TType.MAP) {
            TMap _map = this.readMapBegin();
            HEAD_INFO = new HashMap<>(2 * _map.size);
            for (int i = 0; i < _map.size; ++i) {
                String key = this.readString();
                String value = this.readString();
                HEAD_INFO.put(key, value);
            }
            this.readMapEnd();
        }
        this.readFieldEnd();
        return HEAD_INFO.size() > 0;
    }

    public Map<String, String> getHead() {
        return HEAD_INFO;
    }

    public void markTFramedTransport(TProtocol in) {
        try {
            Field tioInputStream = TIOStreamTransportFieldsCache.getInstance().getTIOInputStream();
            if (tioInputStream == null) {
                return;
            }
            BufferedInputStream inputStream = (BufferedInputStream) tioInputStream.get(in.getTransport());
            inputStream.mark(0);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }


    /*
     * 重置TFramedTransport流，不影响Thrift原有流程
     */
    public void resetTFramedTransport(TProtocol in) {
        try {
            Field tioInputStream = TIOStreamTransportFieldsCache.getInstance().getTIOInputStream();
            if (tioInputStream == null) {
                return;
            }

            // 无法得到TFramedTransport流的输入流
            BufferedInputStream inputStream = (BufferedInputStream) tioInputStream.get(in.getTransport());
            inputStream.reset();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static class TIOStreamTransportFieldsCache {
        private static TIOStreamTransportFieldsCache instance;
        private final Field inputStream_;
        private final String TIOStreamTransport_inputStream_ = "inputStream_";

        private TIOStreamTransportFieldsCache() throws Exception {
            inputStream_ = TIOStreamTransport.class
                    .getDeclaredField(TIOStreamTransport_inputStream_);
            inputStream_.setAccessible(true);
        }

        public static TIOStreamTransportFieldsCache getInstance()
                throws Exception {
            if (instance == null) {
                synchronized (TIOStreamTransportFieldsCache.class) {
                    if (instance == null) {
                        instance = new TIOStreamTransportFieldsCache();
                    }
                }
            }
            return instance;
        }

        public Field getTIOInputStream() {
            return inputStream_;
        }
    }


    public static class Factory implements TProtocolFactory {

        public TProtocol getProtocol(TTransport trans) {
            return new THeaderServerProtocol(trans);
        }
    }


}
